from datasets import load_dataset
from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
import evaluate
import numpy as np

# 1. 加载数据集
dataset = load_dataset("imdb")

# 2. 加载 tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

# 数据预处理函数
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)

tokenized_datasets = dataset.map(preprocess_function, batched=True)

# 3. 划分数据集
train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(20))  # 取2000样本加快训练
eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(10))

# 4. 加载模型（2分类）
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

# 5. 定义评估指标
accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return accuracy.compute(predictions=predictions, references=labels)

# 6. 设置训练参数
training_args = TrainingArguments(
    output_dir="./results",               # 保存模型 checkpoints
    evaluation_strategy="epoch",          # 每个 epoch 结束后做评估
    save_strategy="epoch",                 # 每个 epoch 保存 checkpoint
    learning_rate=20e-5,                    # 学习率
    per_device_train_batch_size=16,         # train batch size
    per_device_eval_batch_size=16,          # eval batch size
    num_train_epochs=3,                    # 总 epoch 数
    weight_decay=0.1,                     # 权重衰减
    logging_dir="./logs",                  # 日志目录
    logging_strategy="steps",              # 日志策略
    logging_steps=100,                      # 每 10 步记录一次
    load_best_model_at_end=True,           # 训练结束加载最佳模型
    metric_for_best_model="accuracy",      # 评估用的最佳指标
)

# 7. 定义 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

# 8. 开始训练
trainer.train()

# 9. 评估模型
eval_results = trainer.evaluate()
print("评估结果：", eval_results)

# 10. 保存最终模型
trainer.save_model("./final_model")